import json
import os
from collections import defaultdict

def coco_to_yolo(coco_file, output_dir):
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 读取COCO格式文件
    with open(coco_file, 'r') as f:
        coco_data = json.load(f)
    
    # 提取类别信息，并创建类别ID映射（COCO ID -> YOLO ID）
    categories = {cat['id']: i for i, cat in enumerate(coco_data['categories'])}
    
    # 创建类别名称文件
    with open(os.path.join(output_dir, 'classes.txt'), 'w') as f:
        for cat in sorted(coco_data['categories'], key=lambda x: categories[x['id']]):
            f.write(f"{cat['name']}\n")
    
    # 获取图像信息
    images = {img['id']: img for img in coco_data['images']}
    
    # 按图像ID组织标注
    image_annotations = defaultdict(list)
    for ann in coco_data['annotations']:
        image_annotations[ann['image_id']].append(ann)
    
    # 转换为YOLO格式
    for img_id, anns in image_annotations.items():
        img_info = images[img_id]
        img_width = img_info['width']
        img_height = img_info['height']
        
        # 为每个图像创建一个标签文件
        filename = os.path.splitext(img_info['file_name'])[0]
        with open(os.path.join(output_dir, f"{filename}.txt"), 'w') as f:
            for ann in anns:
                # 获取边界框信息
                bbox = ann['bbox']  # COCO格式: [x, y, width, height]（x,y是左上角坐标）
                
                # 转换为YOLO格式: [class_id, x_center, y_center, width, height]（归一化）
                x_center = (bbox[0] + bbox[2] / 2) / img_width
                y_center = (bbox[1] + bbox[3] / 2) / img_height
                width = bbox[2] / img_width
                height = bbox[3] / img_height
                
                # COCO ID -> YOLO ID
                class_id = categories[ann['category_id']]
                
                # 写入YOLO格式
                f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
    
    print(f"转换完成！标签文件已保存到 {output_dir} 目录")

if __name__ == "__main__":
    coco_file = r"D:\深度学习课设\instances_test2019.json"
    output_dir = r"D:\深度学习课设\yolo_labels"
    coco_to_yolo(coco_file, output_dir) 